# -*- coding: utf-8 -*-

from flask import Flask, render_template, request, jsonify
from PIL import Image
from auth import auth_login, require_login
from models import *
from oss_utils import upload_stream

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = './'

# 配置允许的文件格式
ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'bmp', 'heic', 'JPG', 'JPEG', 'PNG', 'BMP', 'HEIC'}
MAX_FILE_SIZE_MB = 5
MIN_FILE_SIZE_KB = 5
MIN_RESOLUTION = (150, 150)
MAX_EDGE_LENGTH = 4096


def allowed_file(filename):
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# 更新配置以符合新的规则
SHOE_ALLOWED_EXTENSIONS = {'jpeg', 'jpg', 'png', 'bmp', 'webp', 'avif'}
SHOE_MAX_FILE_SIZE_MB = 5
SHOE_RECOMMENDED_ASPECT_RATIO = (4, 3)
SHOE_MAX_ASPECT_RATIO = (3, 2)


def allowed_shoe_file(filename):
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in SHOE_ALLOWED_EXTENSIONS


def check_shoe_image(file):
    try:
        img = Image.open(file)
        file.seek(0, os.SEEK_END)
        file_size = file.tell()
        file.seek(0)
        if file_size > MAX_FILE_SIZE_MB * 1024 * 1024:
            return False, "文件大小建议不超过5MB"
        width, height = img.size
        aspect_ratio = width / height
        if aspect_ratio > SHOE_MAX_ASPECT_RATIO[0] / SHOE_MAX_ASPECT_RATIO[1]:
            return False, f"图片的长宽比应小于3:2"
        if abs(aspect_ratio - (SHOE_RECOMMENDED_ASPECT_RATIO[0] / SHOE_RECOMMENDED_ASPECT_RATIO[1])) > 0.1:
            return False, f"推荐的长宽比为4:3"
        return True, None
    except Exception as e:
        return False, str(e)


def check_image(file):
    try:
        img = Image.open(file)
        file.seek(0, os.SEEK_END)
        file_size = file.tell()
        file.seek(0)
        if not MIN_FILE_SIZE_KB * 1024 <= file_size <= MAX_FILE_SIZE_MB * 1024 * 1024:
            return False, "文件大小必须在5KB~5MB之间。"
        width, height = img.size
        if width < MIN_RESOLUTION[0] or height < MIN_RESOLUTION[1]:
            return False, "图像分辨率太低。分辨率至少为150x150像素。"
        if max(width, height) > MAX_EDGE_LENGTH:
            return False, "图像的最大边长不能超过4096像素。"
        return True, None
    except Exception as e:
        return False, str(e)


@app.route('/api/upload', methods=['POST'])
@require_login
def upload_file():
    try:
        if 'file' not in request.files:
            return api_response(False, message="未找到文件部分", status_code=409)
        file = request.files['file']
        if file.filename == '':
            return api_response(False, message="没有选择文件", status_code=409)
        if file and allowed_file(file.filename):
            valid, error_msg = check_image(file)
            if not valid:
                return api_response(False, message=error_msg, status_code=400)

            oss_sign_url = upload_stream(file.filename, file.stream)
            return api_response(True, data={"url": oss_sign_url})
        else:
            return api_response(False, message="文件类型不支持", status_code=409)
    except Exception as e:
        return api_response(False, message=f"服务器内部错误: {str(e)}", status_code=500)


@app.route('/api/tryon', methods=['POST'])
@require_login
def handle_tryon():
    """试衣"""
    try:
        data = request.json
        top_garment_url = data.get('top_garment_url')
        person_image_url = data.get('person_image_url')
        bottom_garment_url = data.get('bottom_garment_url', '')

        valid, results = request_tryon_model(top_garment_url, person_image_url, bottom_garment_url)
        if not valid:
            return api_response(False, message=results, status_code=400)

        return api_response(True, data=results)
    except Exception as e:
        return api_response(False, message=f"服务器内部错误: {str(e)}", status_code=500)


@app.route('/api/tryonRefiner', methods=['POST'])
@require_login
def handle_tryon_refiner():
    """试衣精修"""
    try:
        data = request.json
        top_garment_url = data.get('top_garment_url')
        person_image_url = data.get('person_image_url')
        bottom_garment_url = data.get('bottom_garment_url', '')
        coarse_image_url = data.get('coarse_image_url')
        gender = get_gender(coarse_image_url)
        valid, results = request_tryon_refiner(top_garment_url, person_image_url, coarse_image_url, gender,
                                               bottom_garment_url)
        if not valid:
            api_response(False, message=results, status_code=400)

        return api_response(True, data=results)
    except Exception as e:
        return api_response(False, message=f"服务器内部错误: {str(e)}", status_code=500)


@app.route('/api/upload/shoe', methods=['POST'])
@require_login
def upload_shoe_file():
    try:
        if 'file' not in request.files:
            return api_response(False, message="未找到文件部分", status_code=409)
        file = request.files['file']
        if file.filename == '':
            return api_response(False, message="没有选择文件", status_code=409)
        if file and allowed_shoe_file(file.filename):
            valid, error_msg = check_shoe_image(file)
            if not valid:
                api_response(False, message=error_msg, status_code=400)
            oss_sign_url = upload_stream(file.filename, file.stream)
            return api_response(True, data={"url": oss_sign_url})
        else:
            return api_response(False, message="文件类型不支持", status_code=409)
    except Exception as e:
        return api_response(False, message=f"服务器内部错误: {str(e)}", status_code=500)


@app.route('/api/tryShoe', methods=['POST'])
@require_login
def handle_try_shoe():
    """试衣精修"""
    try:
        data = request.json
        template_image_url = data.get('template_image_url')
        shoe_image_url = data.get('shoe_image_url')
        valid, results = request_shoe_model(template_image_url, shoe_image_url)
        if not valid:
            api_response(False, message=results, status_code=400)

        return api_response(True, data=results)
    except Exception as e:
        return api_response(False, message=f"服务器内部错误: {str(e)}", status_code=500)


@app.route('/login', methods=['POST', 'GET'])
def login():
    return auth_login()


@app.route('/', methods=['GET'])
@require_login
def index():
    return render_template('index.html')


def api_response(status, message='', data=None, status_code=200):
    """
    构造API响应的辅助函数。

    :param status: 响应状态（布尔值或字符串：例如 True/False 或 'success', 'error'）
    :param message: 响应消息（可选）
    :param data: 响应数据（可选，默认为空字典）
    :param status_code: HTTP状态码（默认200）
    :return: JSON格式的响应和HTTP状态码
    """
    if isinstance(status, bool):
        status = 'success' if status else 'error'

    response = {
        'status': status,
        'message': message,
        'data': data if data is not None else {}
    }
    return jsonify(response), status_code


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=9000)
